"""
Kokkos-Only Extension Implementation.

This extension wraps pure Kokkos solver implementations for use in GSM.
"""

import json
import os
import subprocess
import logging
from pathlib import Path
from typing import Dict, List, Optional, Any

from gcr_solver_manager.extensions.base_extension import SolverExtension
from gcr_solver_manager.extensions.base_builder import BuilderInterface
from gcr_solver_manager.extensions.base_runner import RunnerInterface
from gcr_solver_manager.extensions.base_validator import ValidatorInterface
from .builder import KokkosOnlyBuilder
from .runner import KokkosOnlyRunner
from .validator import KokkosOnlyValidator

logger = logging.getLogger(__name__)


class KokkosOnlyExtension(SolverExtension):
    """Extension for pure Kokkos solver implementations."""
    
    def __init__(self, source_dir: Optional[Path] = None):
        """Initialize the Kokkos-Only extension."""
        self.source_dir = source_dir or Path("../src/prototypes/kokkos").resolve()
        self._builder = None
        self._runner = None
        self._validator = None
    
    @property
    def name(self) -> str:
        """Extension unique identifier."""
        return "kokkos-only"
    
    @property
    def version(self) -> str:
        """Extension version string."""
        return "1.0.0"
    
    @property
    def supported_frameworks(self) -> List[str]:
        """List of supported framework names."""
        return ["kokkos", "kokkos-only"]
    
    @property
    def supported_solvers(self) -> List[str]:
        """List of supported solver types."""
        return ["gcr", "gmres", "cg", "bicgstab"]
    
    def get_builder(self) -> BuilderInterface:
        """Return builder instance for this extension."""
        if self._builder is None:
            self._builder = KokkosOnlyBuilder(self.source_dir)
        return self._builder
    
    def get_runner(self) -> RunnerInterface:
        """Return runner instance for this extension."""
        if self._runner is None:
            self._runner = KokkosOnlyRunner(self.source_dir)
        return self._runner
    
    def get_validator(self) -> ValidatorInterface:
        """Return validator instance for this extension."""
        if self._validator is None:
            self._validator = KokkosOnlyValidator()
        return self._validator
    
    def get_config_schema(self) -> Dict[str, Any]:
        """Return JSON schema for configuration validation."""
        schema_path = Path(__file__).parent / "config_schema.json"
        if schema_path.exists():
            with open(schema_path, 'r') as f:
                return json.load(f)
        return {}
    
    def can_handle_config(self, config_path: Path) -> bool:
        """Check if extension can handle given configuration."""
        # Check if this is a Kokkos-specific configuration
        
        if not config_path.exists():
            return False
            
        # Check if it's in the Kokkos prototypes directory
        try:
            if 'kokkos' in str(config_path).lower():
                relative_path = config_path.resolve().relative_to(self.source_dir.parent.parent)
                parts = relative_path.parts
                # Handle kokkos directory but exclude hybrid
                if 'kokkos' in parts and 'hybrid' not in parts:
                    return True
        except ValueError:
            pass
        
        # Check file content for Kokkos-specific keywords
        try:
            if config_path.suffix.lower() in ['.xml', '.json', '.yaml']:
                content = config_path.read_text().lower()
                kokkos_keywords = ['kokkos', 'execution_space', 'memory_space', 'parallel_for']
                if any(keyword in content for keyword in kokkos_keywords):
                    # But exclude if it also contains hybrid/petsc keywords
                    exclude_keywords = ['petsc', 'hybrid', 'ksp']
                    if not any(keyword in content for keyword in exclude_keywords):
                        return True
        except Exception:
            pass
        
        return False
    
    def initialize(self, gsm_services: Optional['GSMServices'] = None) -> bool:
        """Initialize extension with GSM services."""
        try:
            # Check for Kokkos installation
            kokkos_dir = os.environ.get('KOKKOS_DIR') or os.environ.get('Kokkos_DIR')
            if not kokkos_dir:
                logger.warning("KOKKOS_DIR not set in environment")
            
            # Validate that Kokkos prototypes directory exists
            if not self.source_dir.exists():
                logger.error(f"Kokkos prototypes directory not found: {self.source_dir}")
                return False
            
            # Check if Kokkos prototype files exist
            kokkos_prototype = self.source_dir / "kokkos_gcr_prototype.cpp"
            if not kokkos_prototype.exists():
                logger.error(f"Kokkos prototype not found: {kokkos_prototype}")
                return False
            
            # Check for CMakeLists.txt
            cmake_file = self.source_dir / "CMakeLists.txt"
            if not cmake_file.exists():
                logger.warning(f"CMakeLists.txt not found: {cmake_file}")
            
            logger.info(f"Kokkos-Only extension initialized with source dir: {self.source_dir}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize Kokkos-Only extension: {e}")
            return False
    
    def cleanup(self) -> bool:
        """Cleanup extension resources."""
        # No specific cleanup needed
        return True